#include <OgreVector3.h>
#include <OgreSceneNode.h>
#include <OgreSceneManager.h>

#include <rviz/ogre_helpers/arrow.h>
#include <rviz/ogre_helpers/billboard_line.h>
#include <rviz/validate_floats.h>

#include <ros/ros.h>

#include <urdf/model.h>
#include "effort_visual.h"

namespace rviz
{
EffortVisual::EffortVisual(Ogre::SceneManager* scene_manager,
                           Ogre::SceneNode* parent_node,
                           boost::shared_ptr<urdf::Model> urdf_model)
{
  scene_manager_ = scene_manager;

  // Ogre::SceneNode s form a tree, with each node storing the
  // transform (position and orientation) of itself relative to its
  // parent.  Ogre does the math of combining those transforms when it
  // is time to render.
  //
  // Here we create a node to store the pose of the Effort's header frame
  // relative to the RViz fixed frame.
  frame_node_ = parent_node->createChildSceneNode();

  //
  urdf_model_ = urdf_model;

  // We create the arrow object within the frame node so that we can
  // set its position and direction relative to its header frame.
  for (std::map<std::string, urdf::JointSharedPtr>::iterator it = urdf_model_->joints_.begin();
       it != urdf_model_->joints_.end(); it++)
  {
    if (it->second->type == urdf::Joint::REVOLUTE)
    {
      std::string joint_name = it->first;
      effort_enabled_[joint_name] = true;
    }
  }
}

EffortVisual::~EffortVisual()
{
  // Delete the arrow to make it disappear.
  for (std::map<std::string, rviz::BillboardLine*>::iterator it = effort_circle_.begin();
       it != effort_circle_.end(); it++)
  {
    delete (it->second);
  }
  for (std::map<std::string, rviz::Arrow*>::iterator it = effort_arrow_.begin();
       it != effort_arrow_.end(); it++)
  {
    delete (it->second);
  }

  // Destroy the frame node since we don't need it anymore.
  scene_manager_->destroySceneNode(frame_node_);
}


void EffortVisual::getRainbowColor(float value, Ogre::ColourValue& color)
{
  value = std::min(value, 1.0f);
  value = std::max(value, 0.0f);

  float h = value * 5.0f + 1.0f;
  int i = floor(h);
  float f = h - i;
  if (!(i & 1))
    f = 1 - f; // if i is even
  float n = 1 - f;

  if (i <= 1)
    color[0] = n, color[1] = 0, color[2] = 1;
  else if (i == 2)
    color[0] = 0, color[1] = n, color[2] = 1;
  else if (i == 3)
    color[0] = 0, color[1] = 1, color[2] = n;
  else if (i == 4)
    color[0] = n, color[1] = 1, color[2] = 0;
  else if (i >= 5)
    color[0] = 1, color[1] = n, color[2] = 0;
}

void EffortVisual::setMessage(const sensor_msgs::JointStateConstPtr& msg)
{
  // for all joints...
  int joint_num = msg->name.size();
  for (int i = 0; i < joint_num; i++)
  {
    std::string joint_name = msg->name[i];
    double effort = msg->effort[i];
    if (!validateFloats(effort))
    {
      ROS_WARN_STREAM("Invalid effort for joint '" << joint_name << "': " << effort);
      continue;
    }
    const urdf::Joint* joint = urdf_model_->getJoint(joint_name).get();
    int joint_type = joint->type;
    if (joint_type == urdf::Joint::REVOLUTE)
    {
      // enable or disable draw
      if (effort_circle_.find(joint_name) != effort_circle_.end() &&
          !effort_enabled_[joint_name]) // enable->disable
      {
        delete (effort_circle_[joint_name]);
        delete (effort_arrow_[joint_name]);
        effort_circle_.erase(joint_name);
        effort_arrow_.erase(joint_name);
      }
      if (effort_circle_.find(joint_name) == effort_circle_.end() &&
          effort_enabled_[joint_name]) // disable -> enable
      {
        effort_circle_[joint_name] = new rviz::BillboardLine(scene_manager_, frame_node_);
        effort_arrow_[joint_name] = new rviz::Arrow(scene_manager_, frame_node_);
      }

      if (!effort_enabled_[joint_name])
        continue;

      // tf::Transform offset = poseFromJoint(joint);
      urdf::JointLimitsSharedPtr limit = joint->limits;
      double max_effort = limit->effort, effort_value = 0.05;

      if (max_effort != 0.0)
      {
        effort_value = std::min(fabs(effort) / max_effort, 1.0) + 0.05;
      }
      else
      {
        effort_value = fabs(effort) + 0.05;
      }

      effort_arrow_[joint_name]->set(0, width_ * 2, width_ * 2 * 1.0, width_ * 2 * 2.0);
      if (effort > 0)
      {
        effort_arrow_[joint_name]->setDirection(orientation_[joint_name] * Ogre::Vector3(-1, 0, 0));
      }
      else
      {
        effort_arrow_[joint_name]->setDirection(orientation_[joint_name] * Ogre::Vector3(1, 0, 0));
      }
      effort_arrow_[joint_name]->setPosition(
          orientation_[joint_name] * Ogre::Vector3(0, 0.05 + effort_value * scale_ * 0.5, 0) +
          position_[joint_name]);
      effort_circle_[joint_name]->clear();
      effort_circle_[joint_name]->setLineWidth(width_);
      for (int i = 0; i < 30; i++)
      {
        Ogre::Vector3 point =
            Ogre::Vector3((0.05 + effort_value * scale_ * 0.5) * sin(i * 2 * M_PI / 32),
                          (0.05 + effort_value * scale_ * 0.5) * cos(i * 2 * M_PI / 32), 0);
        if (effort < 0)
          point.x = -point.x;
        effort_circle_[joint_name]->addPoint(orientation_[joint_name] * point + position_[joint_name]);
      }
      Ogre::ColourValue color;
      getRainbowColor(effort_value, color);
      effort_arrow_[joint_name]->setColor(color.r, color.g, color.b, color.a);
      effort_circle_[joint_name]->setColor(color.r, color.g, color.b, color.a);
    }
  }
}

void EffortVisual::setFrameEnabled(const std::string joint_name, const bool e)
{
  effort_enabled_[joint_name] = e;
}

// Position and orientation are passed through to the SceneNode.
void EffortVisual::setFramePosition(const Ogre::Vector3& position)
{
  frame_node_->setPosition(position);
}

void EffortVisual::setFrameOrientation(const Ogre::Quaternion& orientation)
{
  frame_node_->setOrientation(orientation);
}
// Position and orientation are passed through to the SceneNode.
void EffortVisual::setFramePosition(const std::string joint_name, const Ogre::Vector3& position)
{
  position_[joint_name] = position;
}

void EffortVisual::setFrameOrientation(const std::string joint_name, const Ogre::Quaternion& orientation)
{
  orientation_[joint_name] = orientation;
}

void EffortVisual::setWidth(float w)
{
  width_ = w;
}

void EffortVisual::setScale(float s)
{
  scale_ = s;
}

} // end namespace rviz
